#!/usr/bin/python
import roslib; roslib.load_manifest('speech_recognition')
import rospy
from std_msgs.msg import String
from geometry_msgs.msg import Point


import pyaudio
import subprocess
import xmlrpclib
import sys
import time
import os


LENGTH_IN_MIN = float(sys.argv[1])
SERVER_URL = "http://10.68.0.150:8000"
SERVER = xmlrpclib.Server(SERVER_URL)

CHUNK_SIZE=16000
RATE = 16000
FORMAT = pyaudio.paInt16
DEVICE_NAME = "Kinect"
NUM_CHANNELS = 1
INTERLEAVED = True

################################################################################
# Stream Handling
################################################################################

class StreamManager:
    def start(self):
        with suppress_stdout_stderr():
            self.p = pyaudio.PyAudio()
        device_index = None

        # Find the right audio device
        for device_num in range(self.p.get_device_count()):
            device_info = self.p.get_device_info_by_index(device_num)
            if DEVICE_NAME in device_info['name']:
                device_index = device_num
                self.num_channels = device_info['maxInputChannels']
                print "Using device", device_info['name'],
                print "Number of channels:",self.num_channels

        self.stream = self.p.open(format=FORMAT, channels=self.num_channels,
                                  rate=RATE, input=True, output=True,
                                  input_device_index=device_index,
                                  frames_per_buffer=CHUNK_SIZE)

    def read(self):
        return self.stream.read(CHUNK_SIZE)

    def stop(self):
        self.stream.close()
        self.p.terminate()

import wave
import subprocess
from struct import pack, unpack
from array import array

PATH_STEM = "wavs/recording"

class Recorder:
    def start(self, num_channels, interleaved, sample_width, framerate, chunk_size):
        self.filepaths = []
        self.current_chunks = []
        self.num_channels = num_channels
        self.sample_width = sample_width
        self.rate = framerate
        self.chunk_size = chunk_size
        self.filenum = 0
        self.wav_paths = []
        self.interleaved = interleaved
        
        if self.interleaved == True:
            self.filepaths.append(PATH_STEM)
            self.wav_paths.append([])
        else:
            for i in range(self.num_channels):
                path = PATH_STEM+"_ch"+str(i)
                self.filepaths.append(path)
                self.wav_paths.append([])
        return "server started"

    def close(self):
        self.combine_wavs()
        return "server stopped"
    
    def take_chunk(self, chunk):
        self.process(chunk)
        return "chunk received"

    def process(self, chunk):
        if self.interleaved == True:
            self.current_chunks.append(chunk)
        else:
            data = chunk
            frame = unpack('<' + ('h'*(len(data)/2)), data)
            channels = self.separate_channels(frame)
            for channel in channels:
                data = channel
                channel_chunk = pack('<' + ('h'*len(data)), *data)
                self.current_chunks.append(channel_chunk)
        str_num = str(self.filenum)
        if self.interleaved == True:
            self.save_wav(0)
        else:
            for i in range(self.num_channels):
                self.save_wav(i)
        self.filenum += 1
        
    def save_wav(self, path_index):
        str_num = str(self.filenum)
        path = self.filepaths[path_index]+"_"+str_num+".wav"
        self.wav_paths[path_index].append(path)
        self.write_wav(path, self.current_chunks[-1])
	print "wrote wav", path

    # Takes an unpacked array or list
    # Returns a list of arrays of individual channel data
    def separate_channels(self, frame):
        channels = []
        for i in range(self.num_channels):
            channels.append(array('h'))
        for i in range(len(frame)):
            cur_channel = i % self.num_channels
            channels[cur_channel].append(frame[i])
        return channels

    # Takes packed audio data and adds it to the end of the given wav
    def write_wav(self, path, data):
        wf = wave.open(path, 'wb')
        self.set_wav(wf)
        wf.writeframes(data)
        wf.close()

    def set_wav(self, wf):
        if self.interleaved == False:
            wf.setnchannels(1)
        else:
            wf.setnchannels(self.num_channels)
        wf.setsampwidth(self.sample_width)
        wf.setframerate(self.rate)

    def combine_wavs(self):
        if self.interleaved == True:
            iterations = 1
        else:
            iterations = self.num_channels
        for i in range(iterations):
            data_pieces = []
            for path in self.wav_paths[i]:
                print path
                data = self.wav_data(path)
                data_pieces.append(data)
                delete_file(path)

            combined_data = ""
            for j in range(len(data_pieces)):
                combined_data += data_pieces[j]
            self.write_wav(self.filepaths[i]+".wav", combined_data)
            print "Done writing files!"
            
    def wav_data(self, path):
        wf = wave.open(path, 'rb')
        data = wf.readframes(self.chunk_size)
        wf.close()
        return data
    
def delete_file(path):
    subprocess.call(["rm", path])


class suppress_stdout_stderr(object):
    '''
    A context manager for doing a "deep suppression" of stdout and stderr in 
    Python, i.e. will suppress all print, even if the print originates in a 
    compiled C/Fortran sub-function.
       This will not suppress raised exceptions, since exceptions are printed
    to stderr just before a script exits, and after the context manager has
    exited (at least, I think that is why it lets exceptions through).      

    '''
    def __init__(self):
        # Open a pair of null files
        self.null_fds =  [os.open(os.devnull,os.O_RDWR) for x in range(2)]
        # Save the actual stdout (1) and stderr (2) file descriptors.
        self.save_fds = (os.dup(1), os.dup(2))

    def __enter__(self):
        # Assign the null pointers to stdout and stderr.
        os.dup2(self.null_fds[0],1)
        os.dup2(self.null_fds[1],2)

    def __exit__(self, *_):
        # Re-assign the real stdout/stderr back to (1) and (2)
        os.dup2(self.save_fds[0],1)
        os.dup2(self.save_fds[1],2)
        # Close the null files
        os.close(self.null_fds[0])
        os.close(self.null_fds[1])
    
################################################################################
# Speech Recording
################################################################################

class SpeechRecorder:
    def __init__(self):
        self.sm = StreamManager()
        
    def record(self):
        end_time = time.time()+LENGTH_IN_MIN*60
        self.sm.start()
        self.recorder = Recorder()
        self.recorder.start(self.sm.num_channels, INTERLEAVED,
                            self.sm.p.get_sample_size(FORMAT),
                            RATE, CHUNK_SIZE)
        while time.time() < end_time:
            self.recorder.take_chunk(self.sm.read())

        self.sm.stop()
        self.recorder.close()
    
################################################################################
# Main Loop
################################################################################

if __name__ == '__main__':
    speech_recorder = SpeechRecorder()
    while not rospy.is_shutdown():
        speech_recorder.record()
        sys.exit()
